import utility
import os, scipy, numpy, pickle, cPickle
from Crypto.Cipher import AES
from Crypto.Hash import SHA256
from Crypto.Util import number

class InstructionTable:
    """ Representation of an Instruction Table for a user. """

    """
    - it handles alphas and betas, low-level calculating (alpha,beta or getting x,y from alpha, beta) is done by SecLogin
    - encrypting, decrypting with pwd - does not need to store pwd
    - update rows to replace alpha and beta values
    - It is stored encrypted in memory too.
    """

    def __init__(self, pwd, path, init_tbl=True):
        #init sets and checks path. it doesnt deserialize or set table value, update_rows / deserialize do those.

        #make sure that pwd, path are strings
        try:
            assert(type(pwd) == str)
            assert(type(path) == str)
        except:
            utility.debug('Invalid parameter.')
            return None
        
        self.table = []
        #our table should always have this many rows.
        #this can be checked against when deserializing, to make sure its a proper thing we unpickled.
        self.table_size = len(pwd)

        try:        
            self.home = os.path.abspath(path)
            os.chdir(self.home)
            #If dir exists, file must exist too.
            if (os.path.exists("instTable.tbl") == True):
                self.serialized = True
                utility.debug('Initialized Instruction Table from file.')
            else:
                utility.debug('Instruction Table initialization failed.')
                return None
        except OSError:
            utility.debug('Home directory not found.')
            #this either means its init, or that load failed.
            if (init_tbl):
                try:
                    os.mkdir(self.home)
                    os.chdir(self.home)
                    instTable_file = open("instTable.tbl", 'w')
                    instTable_file.close()
                    self.serialized = False
                    utility.debug('Initialized Instruction Table.')                
                except OSError:
                    utility.debug('Home directory creation failed.')
                    return None
                #IOError, or privilege error
                except:
                    utility.debug('Instruction Table file creation failed.')
                    return None
            else:
                utility.debug('Instruction Table initialization failed on load.')
                return None

    #Private functions, called by serialize/deserialize/update_rows
    """
    Encrypts a parameter of the InstructionTable with the supplied pwd.

    Parameters:
    pwd: the supplied user pwd. It is used lowercase for encryption/decryption.

    Output: True if success, False if error.
    """    
    def _encrypt(self, pwd, parameter):
        try:            
            #get a 32 byte value from pwd
            hashed_pwd = SHA256.new(pwd).digest()
            #use 32 byte value as AES key
            cryptObj = AES.new(hashed_pwd)            
            #encrypt it
            plaintext = number.long_to_bytes(parameter)
    
            #pad it to multiple of 16
            pad_length = cryptObj.block_size - (len(plaintext) % cryptObj.block_size)
            if (pad_length == 0):
                for i in range(cryptObj.block_size):
                    plaintext += chr(0)
            else:
                for i in range(pad_length):
                    plaintext += chr(pad_length)

            ciphertext = cryptObj.encrypt(plaintext)
            return ciphertext
        except:
            raise NameError('Encryption in Instruction Table failed.')
        
    
    """
    Decrypts a parameter of the InstructionTable with the supplied pwd.
    
    Parameters:
    pwd: the supplied user pwd. It is used lowercase for encryption/decryption.

    Output: True if success, False if error.
    """
    def _decrypt(self, pwd, parameter):

        if (type(parameter) != str):
            return None
        else:
            #return number.bytes_to_long(parameter)    
            hashed_pwd = SHA256.new(pwd).digest()
            #use 32 byte value as AES key
            cryptObj = AES.new(hashed_pwd)
            #dec1rypt parameter which is a bytestream ciphertext and turn it into a long
            plaintext = cryptObj.decrypt(parameter)
            pad_length = plaintext[len(plaintext) - 1]

            if (ord(pad_length) == 0):            
                cutoff = len(plaintext) - cryptObj.block_size
            else:
                cutoff = len(plaintext) - ord(pad_length)
            plaintext = plaintext[:cutoff]
            return number.bytes_to_long(plaintext)    

    #Public functions
    """
    Checks whether the InstructionTable is currently in volatile storage or on disk.

    Paremeters: None.
    Returns: the value of self.serialized (True or False).
    """
    def is_serialized(self):
        return self.serialized

    """
    Identical functions to return an alpha or beta value in a row. Decrypts the values on the fly.
    Used during authentication by SecLogin object to poll for alpha, beta values when reconstructing (x,y) pairs. It decrypts the value before returning it.
    
    Parameters:
    index: # of the row of Instruction Table.
    pwd: the pwd used to encryped alpha, beta values.
    Make it all lowercase before usage!

    Returns: alpha or beta
    """
    def get_alpha(self, index, pwd):
        return self._decrypt(pwd, self.table[index][0])
        try:
            #make sure that pwd is a string
            assert(type(pwd) == str)

        #Make sure it does not index self.table in an invalid way.
        except:        
            utility.debug('Invalid index in instruction table. Right here.')
            print self.table
            print self.table[index][0]
            return None

    def get_beta(self, index, pwd):
        try:            
            #make sure that pwd is a string
            assert(type(pwd) == str)
            return self._decrypt(pwd, self.table[index][1])
        #Make sure it does not index self.table in an invalid way.
        except:
            utility.debug('Invalid index in instruction table')
            return None

    """
    Update the InstructionTable by replacing the alpha, beta values. It encrypts the values on the fly.
    SecLogin calls update_rows once it has generated the alpha, beta values. The values

    Parameters:
    rows: values to update the rows with. Its a two dimensional vector (list of alpha, beta pairs)

    Returns: True if update was successful, False otherwise.
    """
    #Replaces rows w/ new values from rows.
    def update_rows(self, rows, pwd):
        #catch potential error raised by _encrypt
        try:
            fn = lambda tpl: map(lambda x: self._encrypt(pwd, x), tpl)
            self.table = map(fn, rows)
            self.table_size = len(self.table)
            return True
            #return map(lambda tpl: map(lambda x: self._encrypt(pwd, x), tpl), rows)
        except NameError:
            return False
        

    """
    Using the supplied pwd, it deserializes the encrypted Instruction Table from disk
    into an in-memory Instruction Table data structure.

    Parameters: None

    Returns: True if deserializtion is successful, False otherwise.
    It returns False if pwd length doesn't fit # of rows in InstructionTable.

    Note: it doesnt decrypt the values with pwd at this point, that is only
    done for the ones that its needed to on, that is the alpha-betas chosen by
    the supplied pwd. This decryption is done in the get_alpha, get_beta functions.
    """
    def deserialize(self):
        try:
            os.chdir(self.home)

            try:
                instTable_file = open("instTable.tbl", 'r')
                try:                    
                    #self.table = pickle.load(instTable_file)
                    self.table = cPickle.load(instTable_file)

                    #check that self.table is proper size, type
                    #It must be a two dimensional array of bytestrings
                    try:
                        assert(type(self.table) == list)
                        assert(len(self.table) == self.table_size)
                        for i in range(len(self.table)):
                            assert(len(self.table[i]) == 2)
                            assert(type(self.table[i][0]) == str)
                            assert(type(self.table[i][1]) == str)
                    #assertion errror
                    except:
                        utility.debug('The unpickled object is invalid.')                       
                        return False
                    
                #except EOFError, ValueError, IndexError, cPickle.UnpickleError:
                except:
                    utility.debug('Pickle error.')
                    instTable_file.close()
                    return False

                instTable_file.close()
            except IOError:
                utility.debug('Instruction Table initialization failed.')
                return False
        except OSError:
            utility.debug('Home directory not found.')
            return False

        self.serialized = False
        #utility.debug('Deserialized Instruction Table.')
        return True

    """
    Using the supplied pwd, it encrypts the alpha, beta values, then serializes the InstructionTable to disk..

    Parameters: None

    Returns: True if serialization is successful, False otherwise.
    """
    def serialize(self):
        try:
            os.chdir(self.home)

            try:
                instTable_file = open("instTable.tbl", 'w')            
                try:                    
                    #pickle.dump(self.table, instTable_file)
                    pObj = cPickle.Pickler(instTable_file)
                    pObj.dump(self.table)
                except:
                    utility.debug('Pickle error.')
                    instTable_file.close()
                    return False                
                instTable_file.close()
            except IOError:
                utility.debug('Instruction Table initialization failed.')
                return False
        except OSError:
            utility.debug('Home directory not found.')
            return False

        self.serialized = True
        return True    

#Insert unit test here for verifying module implementation
if __name__ == "__main__":
    pwd = "passwd"
    insTable = InstructionTable(pwd, "test/", True)
    rows = [[1,2],[3,4]]
    insTable.update_rows(rows, pwd.lower())
    insTable.serialize()
    
    insTable = InstructionTable(pwd, "./", True)
    insTable.update_rows(rows, pwd.lower())
    insTable.serialize()
